# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/11 19:08 
@Author  : ZhangShenao 
@File    : producer_consumer.py 
@Desc    : 使用Coroutine实现生产者-消费者模型
"""
import asyncio
import random
from asyncio import Queue

from bacis.decorator.perf_log_decorator import perf_log


async def producer(queue: Queue, id: int) -> None:
    """
    生产者Coroutine
    :param queue: 异步队列
    :param id: 生产者ID
    """

    # 生产5个随机数
    for i in range(5):
        val = random.randint(1, 100)
        await queue.put(val)
        print(f'[producer-{id}] produce item: {val}')
        await asyncio.sleep(1)


async def consumer(queue: Queue, id: int) -> None:
    """
    消费者Coroutine
    :param queue: 异步队列
    :param id: 消费者ID
    """

    # 开启死循环,从队列中消费数据
    while True:
        val = await queue.get()
        print(f'[consumer-{id}] consume item: {val}')
        await asyncio.sleep(1)


async def start():
    # 创建异步队列
    queue = asyncio.Queue()

    # 创建2个producer任务和2个consumer任务
    producer_tasks = [asyncio.create_task(producer(queue, x)) for x in range(1, 3)]
    consumer_tasks = [asyncio.create_task(consumer(queue, x)) for x in range(1, 3)]

    # 主Coroutine等待
    await asyncio.sleep(10)

    # 取消消费者任务
    for t in consumer_tasks:
        t.cancel()

    # 获取所有任务的异步执行结果
    result = await asyncio.gather(*producer_tasks, *consumer_tasks, return_exceptions=True)

    # 打印执行结果
    # [None, None, CancelledError(''), CancelledError('')]
    print(result)


@perf_log
def main():
    # 开启事件循环,执行异步任务
    asyncio.run(start())


main()
